import numpy as np
import cv2
import pickle
import matplotlib.pyplot as plt
# 加载保存的模型
with open('clf3.pickle', 'rb') as f:
    loaded_clf = pickle.load(f)

# 读取新的图像（这里假设新图像文件名是new_image.tif，你需要根据实际情况修改）
new_image = cv2.imread('Sandstone_2.tif', cv2.IMREAD_GRAYSCALE)
new_image2 = cv2.imread('Sandstone_2_segment.tif', cv2.IMREAD_GRAYSCALE)


# 应用和之前一样的滤波器来处理新图像，获取特征
def apply_filters(image):
    filtered_images = []
    # 均值滤波
    mean_filtered = cv2.blur(image, (5, 5))
    filtered_images.append(mean_filtered)
    
    # 高斯滤波
    gaussian_filtered = cv2.GaussianBlur(image, (5, 5), 0)
    filtered_images.append(gaussian_filtered)
    
    # Sobel滤波
    sobelx = cv2.Sobel(image, cv2.CV_64F, 1, 0, ksize=5)
    sobely = cv2.Sobel(image, cv2.CV_64F, 0, 1, ksize=5)
    sobel_filtered = np.sqrt(sobelx ** 2 + sobely ** 2)
    filtered_images.append(sobel_filtered)
    
    # Canny滤波
    canny_filtered = cv2.Canny(image, 100, 200)
    filtered_images.append(canny_filtered)
    
    return np.stack(filtered_images, axis=-1)

new_filtered_images = apply_filters(new_image)
new_filtered_images2 = apply_filters(new_image2)
# 提取新图像的特征（这里提取方式和训练时一致）

new_features = []
for i in range(new_image.shape[0]):
    for j in range(new_image.shape[1]):
        pixel_features = new_filtered_images[i, j, :]
        new_features.append(pixel_features)

new_X = np.array(new_features)


new_features2 = []
for i in range(new_image2.shape[0]):
    for j in range(new_image2.shape[1]):
        pixel_features2 = new_filtered_images2[i, j, :]
        new_features2.append(pixel_features2)

new_X2 = np.array(new_features2)



# 使用加载的模型进行预测
new_y_pred = loaded_clf.predict(new_X)
new_y_pred2 = loaded_clf.predict(new_X2)

# 将预测结果转换回图像形状
predicted_image = new_y_pred.reshape(new_image.shape)
predicted_image2 = new_y_pred2.reshape(new_image2.shape)



image1_rgb = cv2.cvtColor(predicted_image, cv2.COLOR_BGR2RGB)
image1_rgb2 = cv2.cvtColor(predicted_image2, cv2.COLOR_BGR2RGB)



fig, axs = plt.subplots(1,3)
axs[0].imshow(new_image,cmap='gray')
axs[0].set_title('original Image')
axs[1].imshow(new_image2,cmap='gray')
axs[1].set_title('segment ')
axs[3].imshow(image1_rgb,cmap='gray')
axs[3].set_title('segmentionsion ')

plt.show()

